'''
数据处理脚本 seg  windows
'''
import glob
import json
import os.path
import random
import shutil

import cv2
import numpy as np


class base:
    '''
    labelme格式 to yolo格式
    '''

    def __init__(self):
        pass

    def single_jsontotxt(self, json_path, out_path):
        '''
        seg: json转txt，单个文件，
        :param json_path:输入json文件地址
        :param out_path:输出txt文件地址  cls xywh
        :return:
        '''

        # load json
        with open(json_path, 'r') as load_f:
            content = json.load(load_f)

        # each box
        file_str = ''
        W, H = content['imageWidth'], content['imageHeight']
        for t in content['shapes']:

            # type, (标签名转typeid) 修改它
            # ref = {'yuanzhu_hong': 0, 'yuanzhu_lan': 1, 'fangkuai_hong': 2}
            type = self.ref[t['label']]
            file_str += str(type)
            # xys
            for x, y in t['points']:
                x = x / W
                y = y / H
                file_str += ' ' + str(round(x, 6)) + ' ' + str(round(y, 6))
            file_str += '\n'

            # # cat
            # file_str += str(type) + ' ' + str(round(x, 6)) + ' ' + str(round(y, 6)) + ' ' + str(
            #     round(w, 6)) + ' ' + str(round(h, 6)) + '\n'  # 4月12日更改

        # save
        filename = out_path
        if os.path.exists(filename):
            os.remove(filename)
        # os.mknod(filename)  # win不支持
        # fp = open(filename, mode="r+", encoding="utf-8")
        fp = open(filename, mode="w", encoding="utf-8") # win
        fp.write(file_str[:-1])
        fp.close()

    def make_txt_from_json(self):
        '''
        json to txt
        '''

        ls = glob.glob(self.json)
        total = 0
        for i in ls:
            txt_path = f'{i[:-5]}.txt'
            self.single_jsontotxt(i, txt_path)  # 转换
            total += 1
            print(f'{total}/{int(len(ls))}: saved to {txt_path}')
        print(f'Make_txt_from_json, Done.')

    # def make_yolov5(self):
    #     pass

    def update_yolov5(self):
        '''
        make yolov5 or add a unit to yolov5
        '''

        # source
        self.info()

        # dst
        if not os.path.exists(self.yolov5):
            os.makedirs(self.yolov5)
        out_fir_tra_img = self.yolov5 + 'images/train/'
        out_fir_val_img = self.yolov5 + 'images/val/'
        out_fir_tra_txt = self.yolov5 + 'labels/train/'
        out_fir_val_txt = self.yolov5 + 'labels/val/'
        iltv_dir = [out_fir_tra_img, out_fir_val_img, out_fir_tra_txt, out_fir_val_txt]
        for i in iltv_dir:
            if not os.path.exists(i):
                os.makedirs(i)

        print('before:')
        t_i, v_i, t_tx, v_tx = len(os.listdir(out_fir_tra_img)), len(os.listdir(out_fir_val_img)), \
                               len(os.listdir(out_fir_tra_txt)), len(os.listdir(out_fir_val_txt))
        print(f'{out_fir_tra_img}:{t_i}')
        print(f'{out_fir_val_img}:{v_i}')
        print(f'{out_fir_tra_txt}:{t_tx}')
        print(f'{out_fir_val_txt}:{v_tx}')

        # 划分 各类随机抽1%
        ls_img = glob.glob(self.jpg)
        random.shuffle(ls_img)
        flag = max(int(0.01 * len(ls_img)) + 1, 1)  # 1%
        ls_img_val = ls_img[:flag]
        ls_img_tra = ls_img[flag:]

        # 添加至验证集
        for i in ls_img_val:
            # jpg
            # nam = i.split('/')[-1]
            nam = i.split('\\')[-1] # windows
            shutil.copy(i, out_fir_val_img + f'{self.buff}_{nam}')  # 添加至val_img
            # lab
            # txt_abpath = self.txt + nam[:-4] + '.txt'  #
            txt_abpath = f'{i[:-4]}.txt'  #
            if os.path.exists(txt_abpath):  # 如果有标签文件
                shutil.copy(txt_abpath, out_fir_val_txt + f'{self.buff}_{nam[:-4]}.txt')

        # 添加至训练集
        for ind, i in enumerate(ls_img_tra):
            # jpg
            # nam = i.split('/')[-1]
            nam = i.split('\\')[-1]  # windows
            shutil.copy(i, out_fir_tra_img + f'{self.buff}_{nam}')  # 添加至tra_img
            # lab
            txt_abpath = f'{i[:-4]}.txt'
            if os.path.exists(txt_abpath):
                shutil.copy(txt_abpath, out_fir_tra_txt + f'{self.buff}_{nam[:-4]}.txt')  # 添加至tra_lab

        # check
        print('after:')
        t_i_d, v_i_d, t_tx_d, v_tx_d = len(os.listdir(out_fir_tra_img)), len(os.listdir(out_fir_val_img)), \
                                       len(os.listdir(out_fir_tra_txt)), len(os.listdir(out_fir_val_txt))
        print(f'{out_fir_tra_img}:{t_i_d}')
        print(f'{out_fir_val_img}:{v_i_d}')
        print(f'img added: {t_i_d + v_i_d - t_i - v_i}')
        print(f'{out_fir_tra_txt}:{t_tx_d}')
        print(f'{out_fir_val_txt}:{v_tx_d}')
        print(f'lab added: {t_tx_d + v_tx_d - t_tx - v_tx}')

        print(f'Update_yolov5, Done.')

    def remove_yolov5(self, buff):
        '''
        remove a unit from yolov5
        :return:
        '''
        ls = glob.glob(f'{self.yolov5}*/*/*')
        for ind, i in enumerate(ls):
            # if i.split('/')[-1][:len(buff)] == buff:  # 根据前缀删除
            if i.split('\\')[-1][:len(buff)] == buff:  # windows
                os.remove(i)
                print(f'{ind}/{len(ls)}:{i}')

    def info(self):
        # ls = glob.glob(self.glob_str)
        # print(self.buff)
        # print(
        #     f'origin:{len(ls)}\njpg:{len(os.listdir(self.jpg))}\njson:{len(os.listdir(self.json))}\ntxt:{len(os.listdir(self.txt))}\n')
        pass

    def info_txt(self):
        pass

    def resize_imgs(self, scale = 1.0): # 缩小训练集图片分辨率，加快训练

        imgs_glob = self.yolov5+'images/*/*.jpg'
        ls = glob.glob(imgs_glob)
        for i in ls:
            img = cv2.imread(i)
            img = cv2.resize(img, dsize=None,fx=scale,fy=scale)
            cv2.imwrite(i, img)
        print(f'resized done, total {len(ls)}, scale {scale}')
class labelme2yoloseg(base):
    jpg = r'D:\data\231207huoni\trainseg_all\20231207_simu\*.jpg'
    json = r'D:\data\231207huoni\trainseg_all\20231207_simu\*.json'
    txt = r'D:\data\231207huoni\trainseg_all\20231207_simu\*.txt'
    yolov5 = r'D:\data\231207huoni\trainseg_all\imglabs_yolov5/'  # yolov5数据集地址
    buff = 'd2023_1221'
    ref = {'zhihe': 0, 'zhuji_zheng': 1, 'zhuji_fan': 2, 'dianchi_zheng': 3, 'dianchi_fan': 4, 'xianshu': 5, 'shuomingshu': 6, 'wandai': 7 }



    def run(self):
        self.make_txt_from_json()
        self.update_yolov5()
        # self.info()
        self.resize_imgs(scale=0.3) # (2048, 3072, 3)


if __name__ == '__main__':
    # deepcam_baby().remove_yolov5(buff='d03_03')

    # video2imgs().video2imgs_muti()
    # video2imgs().get_half()

    labelme2yoloseg().run()
